/*37218898_306040569 Lihi_Dery_Luba_Golosman*/
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map.Entry;


public class EMAlgorithm {

	private ArrayList<Document> documents;
	private Document corpus ;
	private int vocabularySize; 
	
	//P(Wk|Xi) - the string is Wk and the Event contains an array for all probabilities P(Wk|Xi) 
	private HashMap<String,Event> wordProb_before = new  HashMap<String,Event>();
	//P(Xi)
	private double[] clusterProb_before = new double[Main.NUM_Clusters];
	
	//P(Wk|Xi) - the string is Wk and the Event contains an array for all probabilities P(Wk|Xi) 
	private HashMap<String,Event> wordProb_current = new  HashMap<String,Event>();
	//P(Xi)
	private double[] clusterProb_current = new double[Main.NUM_Clusters];
	
	//Wti for the E stage [cluster][document]
	private double[][] clusterProbInDocument;
	
	public EMAlgorithm(ArrayList<Document> documents, Document corpus) {
		this.documents = documents;
		this.corpus = corpus;
		clusterProbInDocument =new double[Main.NUM_Clusters][documents.size()];
		analizeCorpus();
	}

	//Analyze corpus and count the words in each document 
	private void analizeCorpus()
	{
		for (Document doc : documents) {
			doc.InitializeEventMap();
		}
		corpus.InitializeEventMap();
		vocabularySize = corpus.getEventMap().size();
	}
	
	
	public void calculateExpectation()
	{
		EMInitialization();
		stageE();
		//TODO 
		//while ()
		{
			wordProb_before = wordProb_current;
			clusterProb_before = clusterProb_current;
			wordProb_current = new  HashMap<String,Event>();
			clusterProb_current = new double[Main.NUM_Clusters];
			
			stageM();
		}
	}
	
	private void EMInitialization()
	{
		Document[] clusters = new Document[Main.NUM_Clusters];

		int i =0;
		//clusters the documents
		for (Document doc : documents) { 
			if (clusters[i%Main.NUM_Clusters] == null)
			         clusters[i%Main.NUM_Clusters]=new Document();
			clusters[i%Main.NUM_Clusters].getWords().addAll(doc.getWords());
			i++;
		}
		// Initialize P(Xi)
		for (int j = 0; j < clusterProb_current.length; j++) {
			clusterProb_current[j] = 1D/Main.NUM_Clusters;
		}
		
		// Initialize P(Wk|Xi)
		for (int j = 0; j < clusters.length; j++) {
			
			clusters[j].InitializeEventMap();
			for (Entry<String, Integer> entry : clusters[j].getEventMap().entrySet()) {
				if (wordProb_current.containsKey(entry.getKey()))
				{
					 wordProb_current.get(entry.getKey()).updateClusterProp(j, (double)entry.getValue()/(double)clusters[j].getNumWords());
				}
				else
				{
					Event event = new Event();
					event.updateClusterProp(j, (double)entry.getValue()/(double)clusters[j].getNumWords());	
					wordProb_current.put(entry.getKey(), event);
				}
			}
		}
		//for debug
		for (Entry<String, Event> entry : wordProb_current.entrySet()) {
			System.out.println(entry.getKey() + "  "+ entry.getValue());
		}
		System.out.println();
	}
	
	private void stageE()
	{//TODO
		
	}
	
	private void stageM()
	{//TODO
		
	}
}
